# %%
import os
import torch

from torch.utils.data import DataLoader

from model.seg_and_class import UNet, classifier, train_model_
from utils import preprocess, MyDataset, MyDataset2, tt, tp, intermediate_result, merge_datasets, \
    classification_results, setup_seed
import matplotlib.pyplot as plt

# %%
BATCH_SIZE = 8
EPOCHS = 1

dev = 'cpu'
if torch.cuda.is_available():
    dev = 'cuda'
print("running on " + dev)

setup_seed(707)

# %%

unet = UNet().to(dev)
loss_fn = torch.nn.MSELoss()
opt = torch.optim.Adam(unet.parameters(), lr=2e-4)

classifier = classifier.to(dev)
loss_fn2 = torch.nn.CrossEntropyLoss()
opt2 = torch.optim.Adam(classifier.parameters(), lr=1e-4)


# %%
def main():
    if not os.path.exists('./segmentation_WBC-master/merged_dataset/'):
        for idx, offset in enumerate([0, 300]):
            if not os.path.exists("./segmentation_WBC-master/preproc_ds{0}/".format(idx + 1)):
                os.mkdir('./segmentation_WBC-master/preproc_ds{0}/'.format(idx + 1))
                preprocess('./segmentation_WBC-master/Class Labels of Dataset {0}.csv'.format(idx + 1),
                           "./segmentation_WBC-master/Dataset {0}/".format(idx + 1),
                           "./segmentation_WBC-master/preproc_ds{0}/".format(idx + 1), offset=offset)

        merge_datasets()
        for i in range(1, 3):
            os.rmdir('./segmentation_WBC-master/preproc_ds{0}/'.format(i))

    if not os.path.exists('output'):
        os.mkdir('output')
    merged_dataset = MyDataset("./segmentation_WBC-master/merged_dataset/",
                               "./segmentation_WBC-master/merged_dataset/merged_labels.csv", tt)

    training_set, test_set = torch.utils.data.random_split(merged_dataset, [300, 100])

    training_loader = DataLoader(training_set, batch_size=BATCH_SIZE, shuffle=False)
    test_loader = DataLoader(test_set, batch_size=BATCH_SIZE, shuffle=False)

    if not os.path.exists('output/seg_model_params.pt'):
        train_model_(EPOCHS, unet, opt,
                     training_loader,
                     test_loader,
                     loss_fn, "output/seg_model_params.pt",
                     ['output/seg_train_hist.pt', 'output/seg_test_hist.pt'],
                     dev=dev)
    else:
        unet.load_state_dict(torch.load("output/seg_model_params.pt", map_location=dev))
    if not os.path.exists('output/inter_result'):
        os.mkdir('output/inter_result')
        os.mkdir('output/inter_result/training_set')
        os.mkdir('output/inter_result/test_set')
        os.mkdir('output/inter_result/gt_train')
        os.mkdir('output/inter_result/gt_test')
        intermediate_result(unet, training_set, 'output/inter_result/training_set/',
                            'output/inter_result/gt_train/',
                            'output/inter_result/training_labels.csv', dev)
        intermediate_result(unet, test_set, 'output/inter_result/test_set/',
                            'output/inter_result/gt_test/',
                            'output/inter_result/test_labels.csv', dev)

    training_cls_set = MyDataset2("./output/inter_result/training_set/",
                                  'output/inter_result/training_labels.csv',
                                  tt)
    test_cls_set = MyDataset2("./output/inter_result/test_set/",
                              'output/inter_result/test_labels.csv',
                              tt)

    training_cls_loader = DataLoader(training_cls_set, batch_size=BATCH_SIZE, shuffle=False)
    test_cls_loader = DataLoader(test_cls_set, batch_size=BATCH_SIZE, shuffle=False)

    if not os.path.exists('output/cls_model_params.pt'):
        train_model_(EPOCHS, classifier, opt2,
                     training_cls_loader,
                     test_cls_loader,
                     loss_fn2,
                     'output/cls_model_params.pt',
                     ['output/cls_train_loss.pt', 'output/cls_test_loss.pt'],
                     dev=dev)
    else:
        classifier.load_state_dict(torch.load("output/cls_model_params.pt", map_location=dev))
    classification_results(classifier, training_cls_set, 'output/tarining_cls_result.csv', dev)
    classification_results(classifier, test_cls_set, 'output/test_cls_result.csv', dev)


# %%
if __name__ == '__main__':
    main()
